package io_test

import (
	"bytes"
	"context"
	"testing"
	"time"

	"github.com/prometheus/prometheus/model/rulefmt"
	"github.com/stretchr/testify/assert"

	"github.com/slok/sloth/internal/log"
	k8stransformpromopv1 "github.com/slok/sloth/internal/plugin/k8stransform/prom_operator_prometheus_rule_v1"
	"github.com/slok/sloth/internal/storage/io"
	"github.com/slok/sloth/pkg/common/model"
	plugink8stransformv1 "github.com/slok/sloth/pkg/prometheus/plugin/k8stransform/v1"
)

func TestIOWriterK8sObjectYAMLRepo(t *testing.T) {
	tests := map[string]struct {
		transformerFunc func() plugink8stransformv1.Plugin
		k8sMeta         model.K8sMeta
		slos            model.PromSLOGroupResult
		expYAML         string
		expErr          bool
	}{
		"Having a multiple SLO alert and recording rules should render correctly.": {
			transformerFunc: func() plugink8stransformv1.Plugin {
				p, _ := k8stransformpromopv1.NewPlugin()
				return p
			},
			k8sMeta: model.K8sMeta{
				Name:        "test-name",
				Namespace:   "test-ns",
				Labels:      map[string]string{"lk1": "lv1"},
				Annotations: map[string]string{"ak1": "av1"},
			},
			slos: model.PromSLOGroupResult{SLOResults: []model.PromSLOResult{
				{
					SLO: model.PromSLO{ID: "testa"},
					PrometheusRules: model.PromSLORules{
						SLIErrorRecRules: model.PromRuleGroup{
							Name: "sloth-slo-sli-recordings-testa",
							Rules: []rulefmt.Rule{
								{
									Record: "test:record-a1",
									Expr:   "test-expr-a1",
									Labels: map[string]string{"test-label": "a-1"},
								},
								{
									Record: "test:record-a2",
									Expr:   "test-expr-a2",
									Labels: map[string]string{"test-label": "a-2"},
								},
							}},
						MetadataRecRules: model.PromRuleGroup{
							Name: "sloth-slo-meta-recordings-testa",
							Rules: []rulefmt.Rule{
								{
									Record: "test:record-a3",
									Expr:   "test-expr-a3",
									Labels: map[string]string{"test-label": "a-3"},
								},
								{
									Record: "test:record-a4",
									Expr:   "test-expr-a4",
									Labels: map[string]string{"test-label": "a-4"},
								},
							}},
						AlertRules: model.PromRuleGroup{
							Name:     "sloth-slo-alerts-testa",
							Interval: 15 * time.Minute, // Custom interval.
							Rules: []rulefmt.Rule{
								{
									Alert:       "testAlertA1",
									Expr:        "test-expr-a1",
									Labels:      map[string]string{"test-label": "a-1"},
									Annotations: map[string]string{"test-annot": "a-1"},
								},
								{
									Alert:       "testAlertA2",
									Expr:        "test-expr-a2",
									Labels:      map[string]string{"test-label": "a-2"},
									Annotations: map[string]string{"test-annot": "a-2"},
								},
							}},
					},
				},
				{
					SLO: model.PromSLO{ID: "testb"},
					PrometheusRules: model.PromSLORules{
						SLIErrorRecRules: model.PromRuleGroup{
							Name: "sloth-slo-sli-recordings-testb",
							Rules: []rulefmt.Rule{
								{
									Record: "test:record-b1",
									Expr:   "test-expr-b1",
									Labels: map[string]string{"test-label": "b-1"},
								},
							}},
						MetadataRecRules: model.PromRuleGroup{
							Name: "sloth-slo-meta-recordings-testb",
							Rules: []rulefmt.Rule{
								{
									Record: "test:record-b2",
									Expr:   "test-expr-b2",
									Labels: map[string]string{"test-label": "b-2"},
								},
							}},
						AlertRules: model.PromRuleGroup{
							Name: "sloth-slo-alerts-testb",
							Rules: []rulefmt.Rule{
								{
									Alert:       "testAlertB1",
									Expr:        "test-expr-b1",
									Labels:      map[string]string{"test-label": "b-1"},
									Annotations: map[string]string{"test-annot": "b-1"},
								},
							}},
						ExtraRules: []model.PromRuleGroup{
							{
								Name:     "sloth-slo-extra-rules-000-testb",
								Interval: 42 * time.Minute,
								Rules: []rulefmt.Rule{
									{
										Alert:       "testAlertZ1",
										Expr:        "test-expr-z1",
										Labels:      map[string]string{"test-label": "z-1"},
										Annotations: map[string]string{"test-annot": "z-1"},
									},
								}},
							{}, // Should be skipped.
							{
								Name: "sloth-slo-extra-rules-001-testb",
								Rules: []rulefmt.Rule{
									{
										Alert:       "testAlertZ2",
										Expr:        "test-expr-z2",
										Labels:      map[string]string{"test-label": "z-2"},
										Annotations: map[string]string{"test-annot": "z-2"},
									},
									{
										Alert:       "testAlertZ3",
										Expr:        "test-expr-z3",
										Labels:      map[string]string{"test-label": "z-3"},
										Annotations: map[string]string{"test-annot": "z-3"},
									},
								},
							},
						},
					},
				},
			}},
			expYAML: `
---
# Code generated by Sloth (dev): https://github.com/slok/sloth.
# DO NOT EDIT.

apiVersion: monitoring.coreos.com/v1
kind: PrometheusRule
metadata:
  annotations:
    ak1: av1
  labels:
    app.kubernetes.io/component: SLO
    app.kubernetes.io/managed-by: sloth
    lk1: lv1
  name: test-name
  namespace: test-ns
spec:
  groups:
  - name: sloth-slo-sli-recordings-testa
    rules:
    - expr: test-expr-a1
      labels:
        test-label: a-1
      record: test:record-a1
    - expr: test-expr-a2
      labels:
        test-label: a-2
      record: test:record-a2
  - name: sloth-slo-meta-recordings-testa
    rules:
    - expr: test-expr-a3
      labels:
        test-label: a-3
      record: test:record-a3
    - expr: test-expr-a4
      labels:
        test-label: a-4
      record: test:record-a4
  - interval: 15m
    name: sloth-slo-alerts-testa
    rules:
    - alert: testAlertA1
      annotations:
        test-annot: a-1
      expr: test-expr-a1
      labels:
        test-label: a-1
    - alert: testAlertA2
      annotations:
        test-annot: a-2
      expr: test-expr-a2
      labels:
        test-label: a-2
  - name: sloth-slo-sli-recordings-testb
    rules:
    - expr: test-expr-b1
      labels:
        test-label: b-1
      record: test:record-b1
  - name: sloth-slo-meta-recordings-testb
    rules:
    - expr: test-expr-b2
      labels:
        test-label: b-2
      record: test:record-b2
  - name: sloth-slo-alerts-testb
    rules:
    - alert: testAlertB1
      annotations:
        test-annot: b-1
      expr: test-expr-b1
      labels:
        test-label: b-1
  - interval: 42m
    name: sloth-slo-extra-rules-000-testb
    rules:
    - alert: testAlertZ1
      annotations:
        test-annot: z-1
      expr: test-expr-z1
      labels:
        test-label: z-1
  - name: sloth-slo-extra-rules-001-testb
    rules:
    - alert: testAlertZ2
      annotations:
        test-annot: z-2
      expr: test-expr-z2
      labels:
        test-label: z-2
    - alert: testAlertZ3
      annotations:
        test-annot: z-3
      expr: test-expr-z3
      labels:
        test-label: z-3
`,
		},
	}

	for name, test := range tests {
		t.Run(name, func(t *testing.T) {
			assert := assert.New(t)

			var gotYAML bytes.Buffer
			repo := io.NewIOWriterK8sObjectYAMLRepo(&gotYAML, test.transformerFunc(), log.Noop)
			err := repo.StoreSLOs(context.TODO(), test.k8sMeta, test.slos)

			if test.expErr {
				assert.Error(err)
			} else if assert.NoError(err) {
				assert.Equal(test.expYAML, gotYAML.String())
			}
		})
	}
}
